package svr

import (
	"context"
	"errors"
	"gitee.com/yysf-xin/dproxy/pkg/db"
	"gitee.com/yysf-xin/dproxy/pkg/tools"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/net/proxy"
	"golang.org/x/sync/singleflight"
	"io/ioutil"
	"net"
	"net/url"
	"os"
	"strconv"
	"sync"
	"time"
)

type SshClient struct {
	sshCfg *ssh.ClientConfig
	Cfg    *db.SshServer
	conn   *ssh.Client
	url    *url.URL
	l      sync.RWMutex
	group  singleflight.Group
	sshGroup  singleflight.Group

	Dialer proxy.Dialer
	Name   string
}

func (c *SshClient) String() string {
	return c.Name
}

func (c *SshClient) Start() (err error) {
	_, err, _ = c.group.Do(strconv.Itoa(int(c.Cfg.Id)), c.doConnect)
	return
}

func (c *SshClient) Stop() {
	if c.conn != nil {
		err := c.conn.Close()
		if err != nil {
			tools.Log.Println(c, "closed err", err)
		}
		c.conn = nil
	}
	c.group.Forget(strconv.Itoa(int(c.Cfg.Id)))
}

func (c *SshClient) Close() error {
	if c.conn == nil {
		return nil
	}
	return c.conn.Close()
}

func NewSshClient(cfg *db.SshServer) (c *SshClient, err error) {
	var sshCfg ssh.ClientConfig
	sshCfg.HostKeyCallback = ssh.InsecureIgnoreHostKey()
	remote, err := url.Parse(cfg.Remote)
	if err != nil {
		return
	}
	sshCfg.User = remote.User.Username()
	sshCfg.Timeout = 15 * time.Second
	// 1) try RSA keyring first
	if cfg.IdRsa != "" {
		var signer ssh.Signer
		if signer, err = ssh.ParsePrivateKey([]byte(cfg.IdRsa)); err != nil {
			var pem []byte
			if pem, err = ioutil.ReadFile(cfg.IdRsa); err == nil {
				signer, err = ssh.ParsePrivateKey(pem)
			}
		}
		if signer != nil {
			sshCfg.Auth = append(sshCfg.Auth, ssh.PublicKeys(signer))
		}
	}

	// 2) try password
	if pass, ok := remote.User.Password(); ok {
		sshCfg.Auth = append(sshCfg.Auth, ssh.Password(pass))
	}

	// 3) try ssh agent
	if sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK")); err == nil {
		agentAuth := ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers)
		sshCfg.Auth = append(sshCfg.Auth, agentAuth)
	}

	if len(sshCfg.Auth) == 0 {
		err = errors.New("invalid auth method, please add password or generate ssh keys")
		return
	}
	c = &SshClient{
		sshCfg: &sshCfg,
		Cfg:    cfg,
		url:    remote,
		Name:   strconv.Itoa(int(cfg.Id)) + "@SSH",
	}
	return
}

var ERR = errors.New("not connect")

func checkConnect(cli *ssh.Client) (err error) {
	if cli.Conn == nil {
		return ERR
	}
	_, _, err = cli.Conn.SendRequest("keepalive@golang.org", true, nil)
	return
}
func (c *SshClient) Dial(network, addr string) (net.Conn, error) {
	return c.DialContext(context.Background(), network, addr)
}
func (c *SshClient) DialContext(ctx context.Context, network, addr string) (conn net.Conn, err error) {
	c.l.RLock()
	cli := c.conn
	c.l.RUnlock()
	if cli != nil {
		if conn, err = cli.Dial(network, addr); err == nil {
			return
		}
		if _, ok := err.(*ssh.OpenChannelError); ok {
			return
		}

	}
	v, err, _ := c.group.Do(addr, c.doConnect)
	if err == nil {
		if v != nil {
			cli = v.(*ssh.Client)
			return cli.Dial(network, addr)
		}
	}
	err = net.ErrClosed
	return
}

func (c *SshClient) doConnect() (v interface{}, err error) {
	cli := c.conn
	if cli != nil {
		if err = checkConnect(cli); err != nil {
			tools.Log.Println("ssh dial err",c.url.Host,err)
			err = cli.Close()
			c.l.Lock()
			c.conn = nil
			c.l.Unlock()
		} else {
			return nil, nil
		}
		tools.Log.Println("reconnect ssh",c.url.Host)
	}
	v,err,_ = c.sshGroup.Do(c.url.Host, func() (v interface{}, err error)  {
		if c.Dialer != nil {
			var conn net.Conn
			conn, err = c.Dialer.Dial("tcp", c.url.Host)
			if err != nil {
				return
			}
			sshConn, chans, reqs, err := ssh.NewClientConn(conn, c.url.Host, c.sshCfg)
			if err != nil {
				return nil, err
			}
			cli = ssh.NewClient(sshConn, chans, reqs)
		} else {
			cli, err = ssh.Dial("tcp", c.url.Host, c.sshCfg)
		}

		if cli != nil {
			c.l.Lock()
			c.conn = cli
			c.l.Unlock()
		}
		v = cli
		return
	})
	return v, err
}
